import torch
N,D_in,H,D_out = 64,1000,100,10

#随机创建一些训练数据
x = torch.randn(N,D_in)
y = torch.randn(N,D_out)

model = torch.nn.Sequential(
    torch.nn.Linear(D_in,H),
    torch.nn.ReLU(),
    torch.nn.Linear(H,D_out)
)
# model = model.to("cuda:0")

torch.nn.init.normal_(model[0].weight)
torch.nn.init.normal_(model[2].weight)

loss_fn = torch.nn.MSELoss(reduction = 'sum')

# learning_rate = 1e-3
# optimizer = torch.optim.Adam(model.parameters(),lr=learning_rate)

learning_rate = 1e-6
optimizer = torch.optim.SGD(model.parameters(),lr=learning_rate)


for it in range(500):
    # Forward pass
    y_pred = model(x)

    # compute loss
    loss = loss_fn(y_pred, y)
    print(it,loss.item())

    optimizer.zero_grad()

    # Backward pass
    loss.backward()

    # update weights pf w1 and w2
    optimizer.step()